# =============================================================================
# Functions to split a SOFIA HAWC+ fits file into different extensions
# =============================================================================
from astropy.io import fits

# =============================================================================
# Function to split a HAWC+ fits cube into its main components
# =============================================================================
def SplitHAWC(SourceName, FitsLocation, OutFolder, Coord):
    # Loading HAWC+ fits cube
    CubeHAWC = fits.open(FitsLocation)
    # Extension names and numbers to be used
    ExtNum = [0, 1, 2, 3, 4, 5, 8, 9, 10, 11, 12, 13, 14] # Numbers for the relevant HAWC+ extensions
    ExtNames = ['I', 'dI', 'Q', 'dQ', 'U', 'dU', 'P', 'dP', 'O', 'B', 'dO', 'PI', 'dPI'] # Related extension names
    # Creating a new fits file for each extension
    ExtSize = len(ExtNum)
    print('Dividing the HAWC+ cube for '+SourceName)
    for i in range (0,ExtSize):
        # Creating a new fits container from the given extension
        ExtFits = fits.PrimaryHDU(data=CubeHAWC[ExtNum[i]].data, header=CubeHAWC[ExtNum[i]].header)
        # Saving the new fits file
        ExtFileName = OutFolder+SourceName+'_'+Coord+'_'+ExtNames[i]+'.fits' # Name for each new fits file
        ExtFits.writeto(ExtFileName, overwrite=True) # Writing each new fits files


# =============================================================================
# Function to split a Planck fits cube into its main components
# =============================================================================
def SplitPlanck(SourceName, FitsLocation, OutFolder):
    # Loading Planck cube
    CubePlanck = fits.open(FitsLocation)
    # Extension names
    ExtNames = ['I', 'Q', 'U', 'PI', 'P', 'O', 'B'] # Related extension names
    # Creating a new fits file for each extension
    ExtSize = len(ExtNames)
    print('Dividing the HAWC+ cube for '+SourceName)
    for i in range (0,ExtSize):
        # Creating a new fits container from the given extension
        ExtFits = fits.PrimaryHDU(data=CubePlanck[i].data, header=CubePlanck[i].header)
        # Saving the new fits file
        ExtFileName = OutFolder+SourceName+'_Planck_'+ExtNames[i]+'.fits' # Name for each new fits file
        ExtFits.writeto(ExtFileName, overwrite=True) # Writing each new fits files